#include "crossbit/Exit.h"
#include "arch/x86/X86TraceManager.h"
#include "iostream"
using namespace crossbit;


X86TraceManager::X86TraceManager( MemoryManager *mm ) : TraceManager(mm)
{
    /***********************************************
     * Machine code: 8d8a00000000 e305 e900000000 
     * e900000000 90909090
     * asm:    lea (%edx), %ecx  
     *         jecxz label   // it could just jump short if ecx == 0
     *         jmp 00000000
     * lable:  jmp 00000000
     *
     *********************************************/
    *((XTInt32*)&sample[0]) = 0x00008a8d;
    *((XTInt32*)&sample[4]) = 0x05e30000;
    *((XTInt32*)&sample[8]) = 0x000000e9;
    *((XTInt32*)&sample[12]) = 0x0000e900;
    *((XTInt32*)&sample[16]) = 0x90900000;

}

X86TraceManager::~X86TraceManager()
{
    //empty now
}
void X86TraceManager::linkSuperBlock( TBlock* former_tblock , TBlock* super_tblock ){
    XTMemAddr entry = super_tblock -> getSPC();
    std::map<XTMemAddr , std::vector<XTMemAddr>* >::iterator it = GLOBAL::return_map.find(entry);
    if( it != GLOBAL::return_map.end() ) {
        vector<XTMemAddr> *vt = (*it).second;
        for( XTInt32 i = 0 ; i < vt -> size() ; i++ ) {
            *((XTMemAddr*) vt -> at(i)) = super_tblock -> enterAddr();
        }
        delete vt;
        GLOBAL::return_map.erase(it);
    }

    XTUint8* t_addr = (XTUint8 *)former_tblock -> enterAddr();

    *t_addr  = 0xe9;  //jmp disp
    t_addr++;
    *((XTInt32 *)t_addr) = super_tblock -> enterAddr() - former_tblock ->enterAddr() - 5;

    XTMemAddr addr = former_tblock -> enterAddr();
    XTInt32  dis = super_tblock -> enterAddr() - former_tblock -> enterAddr();
    if( link_table.find(addr) != link_table.end() ){
        map<XTMemAddr , vector<XTMemAddr> *>::iterator it = link_table.find( addr );
        vector<XTMemAddr> *vxt = (*it).second;
        for( int i = 0 ; i < vxt -> size() ; i++ ) {
            XTMemAddr jmp_stub = vxt -> at(i);
            *(XTInt32*)jmp_stub += dis;
        }	
        delete	(vxt);
        link_table.erase(it);
    }
    XTMemAddr SPC = former_tblock -> getSPC();
    XTUint32 index = (SPC & GLOBAL::TABLE_MASK) >> 1;
    if( GLOBAL::indirect_hashtable[ index ] == SPC ){
        GLOBAL::indirect_hashtable[ index + 1] = super_tblock -> enterAddr();
    }
}
void X86TraceManager::directLinkTo( TBlock* preBlock, XTInt32 exit, TBlock* block)
{
    Exit* e = GLOBAL::getExit(exit);
    e -> dest = block->enterAddr();

    if( e -> linked == 0 )
    {
        e -> linked = 1;

        XTInt8* p = e -> linkCode;
        *(XTInt32*)p = (XTInt32)( (e -> dest) - (XTInt32)p - 4);
        if( link_table.find( e -> dest ) == link_table.end() ){
            vector<XTMemAddr>* vxt = new vector<XTMemAddr>();
            link_table[e -> dest] = vxt;
        }
        map<XTMemAddr , vector<XTMemAddr> *>::iterator it = link_table.find( e -> dest );
        (*it).second -> push_back((XTMemAddr)p);
    }
}

void X86TraceManager::callLinkTo( TBlock* preBlock, XTInt32 exit, TBlock* block)
{
    Exit* e = GLOBAL::getExit(exit);
    e -> dest = block->enterAddr();

    if( e -> linked == 0 )
    {
        e -> linked = 1;

        XTMemAddr p = (XTMemAddr)(e -> linkCall);
        *(XTInt8*)p = 0xe9;
        p++;
        *(XTInt32*)p = (XTInt32)( (e -> dest) - (XTInt32)p - 4);
        if( link_table.find( e -> dest ) == link_table.end() ){
            vector<XTMemAddr>* vxt = new vector<XTMemAddr>();
            link_table[e -> dest] = vxt;
        }
        map<XTMemAddr , vector<XTMemAddr> *>::iterator it = link_table.find( e -> dest );
        (*it).second -> push_back((XTMemAddr)p);
    }
}


void X86TraceManager::indirectLinkTo( TBlock* preBlock, XTInt32 exit, TBlock* block )
{
    Exit* e = GLOBAL::getExit(exit);
    XTMemAddr SPC = block->getSPC();
    XTMemAddr enter = block->enterAddr();

    if( e -> linked > 5 ) {
        XTUint32 index = (SPC & GLOBAL::TABLE_MASK) >> 1;
        GLOBAL::indirect_hashtable[ index ] = SPC;
        GLOBAL::indirect_hashtable[ index + 1] = enter;
        return ;
    }

    XTMemAddr p;
    if( e -> linked == 0 ) {
        // inline one comparison bucket.
        p = (XTMemAddr) ( e -> linkCode ) - 1 ; // to minus 1 : please refer to encodeJMP() in X86Translator.cpp
        *((XTInt8*)p) = 0x81;	   	
        *((XTInt8*)(p+1)) = 0xFA;  // cmp edx , SPC
        *((XTInt32*)(p+2)) = SPC;
        *((XTInt8*)(p+6)) = 0x0f;
        *((XTInt8*)(p+7)) = 0x84;
        *((XTInt32*)(p+8)) = enter - (p + 12);  // je enter 
        *((XTInt8*)(p+12)) = 0xe9;  // jmp    
        *((XTInt32*)(p+13)) = ( e -> dest - ( p + 17 ));   //  RET 
        e -> linkCode = (XTInt8*)(p+13); 
        e -> linked = 1;

    } else {
        if( e -> linked == 1 ) {
            p = allocatePredictionCodeBlock();
            if ( !p ) {
                return;
            }
            p += 60	; // 5 comparison-branch pairs
            e -> dest = p;
            XTMemAddr temp = p;
            *((XTInt8*)temp) = 0xc7;  // mov exit -> ebx
            temp++;
            *((XTInt8*)temp) = 0xc0 | 3;
            temp++;
            *((XTInt32*)temp) = exit;
            temp += 4;
            *((XTInt8*)temp) = 0xe9;  // jmp to the hashtable_dispatch
            temp++;
            *((XTInt32*)temp) = GLOBAL::hashtable_dispatch - ( temp + 4 );	
            //	*((XTInt32*)temp) = e -> dest  - ( temp + 4 );	
        }
        e -> linked ++;
        e -> dest -= 12;
        p = e -> dest;
        *((XTInt8*)p) = 0x81;	   	
        *((XTInt8*)(p+1)) = 0xFA;  // cmp edx , SPC
        *((XTInt32*)(p+2)) = SPC;
        *((XTInt8*)(p+6)) = 0x0f;
        *((XTInt8*)(p+7)) = 0x84;
        *((XTInt32*)(p+8)) = enter - (p + 12);  // je enter 

        *((XTInt32*)(e -> linkCode )) = p - ( XTMemAddr) ( e -> linkCode ) - 4;

    }	

    if( link_table.find( enter ) == link_table.end() ){
        vector<XTMemAddr>* vxt = new vector<XTMemAddr>();
        link_table[enter] = vxt;
    }
    map<XTMemAddr , vector<XTMemAddr> *>::iterator it = link_table.find( enter );
    (*it).second -> push_back(p + 8);
}

void X86TraceManager::fillPredictionCodeBlock( TBlock* block )
{
    //empty now
}

XTMemAddr X86TraceManager::allocatePredictionCodeBlock()
{
    XTMemAddr buf;

    if( PREDICTIONCODEBLOCKLEN <= tm_size )
    { 
        buf = tm_pool;
        tm_pool += PREDICTIONCODEBLOCKLEN;
        tm_size -= PREDICTIONCODEBLOCKLEN;
    }
    else
    {
        //it's temporarily now
        //allocate new pool here
        std::cerr << "Memory is not enough" << std::endl;
        exit(0);
    }

    return buf;
}
